from abc import ABC, abstractmethod
from typing import Dict, Any, List, Optional
from datetime import datetime
import json
from zhipuai import ZhipuAI
from config import settings
from utils import logger

class BaseAgent(ABC):
    """AI代理基类"""
    
    def __init__(self, name: str, role: str, model: str = None):
        self.name = name
        self.role = role
        self.model = model or settings.default_model
        
        # 使用智谱清言官方SDK
        self.client = self._create_client()
        
        self.conversation_history: List[Dict[str, Any]] = []
        self.metadata = {
            'created_at': datetime.now().isoformat(),
            'total_calls': 0,
            'total_tokens': 0,
            'success_rate': 0.0,
            'model_provider': '智谱清言',
            'model_name': self.model
        }
        
        logger.info(f"🤖 初始化代理: {self.name} - 使用模型: {self.model} (智谱清言)")
    
    def _create_client(self):
        """创建智谱AI客户端"""
        try:
            return ZhipuAI(api_key=settings.zhipu_api_key)
        except Exception as e:
            logger.error(f"智谱AI客户端初始化失败: {e}")
            raise RuntimeError(f"无法初始化智谱AI客户端: {e}")
    
    def _get_system_prompt(self) -> str:
        """获取系统提示词，子类需要实现"""
        return f"""你是一个专业的{self.role}，具有丰富的经验和深厚的技术背景。
你的任务是帮助用户完成{self.name}相关的工作。
请始终保持专业、准确、有用的回复。
用中文回复。"""
    
    def _call_llm(self, prompt: str, context: Optional[str] = None) -> str:
        """调用智谱清言大语言模型"""
        try:
            messages = [{"role": "system", "content": self._get_system_prompt()}]
            
            # 添加上下文
            if context:
                messages.append({"role": "user", "content": f"上下文信息：\n{context}"})
            
            # 添加当前提示
            messages.append({"role": "user", "content": prompt})
            
            # 调用智谱清言API
            response = self.client.chat.completions.create(
                model=self.model,
                messages=messages,
                temperature=settings.temperature,
                max_tokens=settings.max_tokens
            )
            
            # 获取响应内容
            response_content = response.choices[0].message.content
            
            # 更新统计信息
            self.metadata['total_calls'] += 1
            self.metadata['total_tokens'] += len(response_content)
            
            # 记录对话历史
            self.conversation_history.append({
                'timestamp': datetime.now().isoformat(),
                'prompt': prompt,
                'response': response_content,
                'success': True,
                'model': self.model
            })
            
            logger.info(f"🤖 {self.name} 成功响应 ({self.model}, 长度: {len(response_content)})")
            return response_content
            
        except Exception as e:
            error_msg = f"❌ {self.name} 调用{self.model}失败: {str(e)}"
            logger.error(error_msg)
            
            # 记录失败
            self.conversation_history.append({
                'timestamp': datetime.now().isoformat(),
                'prompt': prompt,
                'response': None,
                'error': str(e),
                'success': False,
                'model': self.model
            })
            
            return f"抱歉，我遇到了技术问题：{str(e)}"
    
    @abstractmethod
    def process(self, input_data: Dict[str, Any]) -> Dict[str, Any]:
        """处理输入数据，返回结果"""
        pass
    
    def get_stats(self) -> Dict[str, Any]:
        """获取代理统计信息"""
        successful_calls = sum(1 for h in self.conversation_history if h.get('success', False))
        total_calls = len(self.conversation_history)
        
        self.metadata['success_rate'] = (successful_calls / total_calls * 100) if total_calls > 0 else 0

        return {
            'name': self.name,
            'role': self.role,
            'performance': {
                'total_calls': total_calls,
                'success_rate': f"{(successful_calls / total_calls * 100):.1f}%" if total_calls > 0 else "0%"
            },
            'metadata': {
                'total_tokens': self.metadata['total_tokens']
            }
        }

    def reset_history(self):
        """重置对话历史"""
        self.conversation_history = []
        logger.info(f"🔄 {self.name} 对话历史已重置")
    
    def export_history(self, filepath: str):
        """导出对话历史"""
        try:
            with open(filepath, 'w', encoding='utf-8') as f:
                json.dump({
                    'agent_name': self.name,
                    'agent_role': self.role,
                    'history': self.conversation_history,
                    'stats': self.get_stats()
                }, f, ensure_ascii=False, indent=2)
            logger.info(f"📁 {self.name} 对话历史已导出: {filepath}")
        except Exception as e:
            logger.error(f"❌ 导出失败: {e}")

class AgentManager:
    """代理管理器"""
    
    def __init__(self):
        self.agents: Dict[str, BaseAgent] = {}
        self.execution_log: List[Dict[str, Any]] = []
    
    def register_agent(self, agent: BaseAgent):
        """注册代理"""
        self.agents[agent.name] = agent
        logger.info(f"✅ 代理已注册: {agent.name} ({agent.role})")
    
    def get_agent(self, name: str) -> Optional[BaseAgent]:
        """获取代理"""
        return self.agents.get(name)
    
    def list_agents(self) -> List[str]:
        """列出所有代理"""
        return list(self.agents.keys())
    
    def get_all_stats(self) -> Dict[str, Any]:
        """获取所有代理的统计信息"""
        return {name: agent.get_stats() for name, agent in self.agents.items()}
    
    def execute_workflow(self, workflow_steps: List[Dict[str, Any]], initial_data: Dict[str, Any]) -> Dict[str, Any]:
        """执行工作流"""
        current_data = initial_data.copy()
        results = {}
        
        for step in workflow_steps:
            agent_name = step['agent']
            agent = self.get_agent(agent_name)
            
            if not agent:
                logger.error(f"❌ 代理未找到: {agent_name}")
                continue
            
            try:
                logger.info(f"🚀 执行步骤: {agent_name}")
                result = agent.process(current_data)
                results[agent_name] = result
                
                # 更新数据流
                if isinstance(result, dict):
                    current_data.update(result)
                
                # 记录执行日志
                self.execution_log.append({
                    'timestamp': datetime.now().isoformat(),
                    'agent': agent_name,
                    'success': True,
                    'result_keys': list(result.keys()) if isinstance(result, dict) else ['result']
                })
                
            except Exception as e:
                logger.error(f"❌ 执行失败 {agent_name}: {e}")
                self.execution_log.append({
                    'timestamp': datetime.now().isoformat(),
                    'agent': agent_name,
                    'success': False,
                    'error': str(e)
                })
        
        return results 